# -*- coding: utf-8 -*-
"""
Created on Fri Aug 13 10:02:10 2021

@author: weixifei
"""
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from model_comfeature import Model

from train import train
import torch
from load_data import GAMMA_sub1_dataset
import torch.nn as nn
from utils import seed_torch 

seed_torch()
batchsize = 8 # 批大小,
oct_img_size = [512, 512]
image_size = 512
num_epoch=100 # 迭代次数
val_ratio = 0.2 # 训练/验证数据划分比例，80 / 20
test_root = "../val_data/val_data/multi-modality_images"
num_workers = 0
init_lr = 0.0001

trainset_root='../GAMMA_training data/training_data/multi-modality_images'

img_train_transforms = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomResizedCrop(image_size, scale=(0.8, 1)),
        transforms.Resize((image_size, image_size)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.ColorJitter(contrast=0.4,brightness=0.4),

        # transforms.ColorJitter(contrast=0.4,brightness=0.4),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])

oct_train_transforms = transforms.Compose([
#        transforms.ToPILImage(),
#        transforms.Resize((image_size, image_size)),
#        transforms.RandomHorizontalFlip(),
#        transforms.RandomRotation(15),
#        transforms.ColorJitter(contrast=0.2,brightness=0.2),
        transforms.ToTensor(),
#        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])

img_val_transforms = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((image_size, image_size)),
#        transforms.RandomHorizontalFlip(),
#        transforms.RandomRotation(15),
#        transforms.ColorJitter(contrast=0.2,brightness=0.2),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])

oct_val_transforms =transforms.Compose([
#        transforms.RandomHorizontalFlip(),
#        transforms.RandomRotation(15),
#        transforms.ColorJitter(contrast=0.2,brightness=0.2),
        transforms.ToTensor(),
#        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])
seed_torch(seed=1995)

train_dataset = GAMMA_sub1_dataset(dataset_root=trainset_root, # 训练数据和val数据在文件夹进行分配
                        img_transforms=img_train_transforms,
                        oct_transforms=oct_train_transforms,
                        label_file='../GAMMA_training data/training_data/glaucoma_grading_training_GT.xlsx',
                        mode='train')

val_dataset = GAMMA_sub1_dataset(dataset_root=trainset_root,
                        img_transforms=img_val_transforms,
                        oct_transforms=oct_val_transforms,
                        label_file='../GAMMA_training data/training_data/glaucoma_grading_training_GT.xlsx',
                        mode='val')

train_loader = DataLoader(train_dataset, batch_size=batchsize,pin_memory=True, num_workers=0,shuffle=False)
val_loader = DataLoader(val_dataset, batch_size=1,pin_memory=True, num_workers=0,shuffle=False)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Model().to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)
criterion = nn.CrossEntropyLoss()
train(model, num_epoch, train_loader, val_loader, optimizer, criterion)

